package util.bio.seq;
import java.io.*;
import java.util.*;
import java.util.regex.*;
import util.gen.*;
import util.bio.parsers.*;
import util.bio.annotation.*;

/**Generates a fasta file for each gene model containing exon seqs plus splice junctions. Reverse complements neg strand annotations.
 * Expands exon size by sequenceLengthRadius.*/
public class MakeTranscriptome {

	//fields
	private HashMap<String,File> fastaFiles;
	private int sequenceLengthRadius = 0;
	private HashMap<String,UCSCGeneLine[]> chromSpecificGeneLines;
	private File resultsFastaFile;
	private boolean appendCoordinates = false;
	private boolean makeTranscriptChromosome = false;
	private String chr;
	private String seq;
	private UCSCGeneLine[] models;
	private PrintWriter outSeq = null;
	private PrintWriter outBed = null;
	private int base = 0;


	public MakeTranscriptome(String[] args){
		//process user args and load data
		processArgs(args);
		try {
			outSeq = new PrintWriter (new FileWriter (resultsFastaFile));
			outSeq.println(">chrT");
			if (makeTranscriptChromosome){
				String trimmed = Misc.removeExtension(resultsFastaFile.getCanonicalPath()) +".bed";
				outBed = new PrintWriter (new FileWriter (new File(trimmed)));
			}
		} catch (Exception e){
			e.printStackTrace();
			System.exit(1);
		}

		//for each chromosome of gene models
		System.out.println("Processing...");
		Iterator<String> it = chromSpecificGeneLines.keySet().iterator();
		while (it.hasNext()){
			chr = it.next();
			System.out.println("\n\t" + chr);

			//fetch the sequence
			MultiFastaParser mfp = new MultiFastaParser(fastaFiles.get(chr));
			seq = mfp.getSeqs()[0];

			//for each gene line
			models = chromSpecificGeneLines.get(chr);

			if (makeTranscriptChromosome) makeChrT();
			else makeExons();
		}
		outSeq.close();
		if (makeTranscriptChromosome) outBed.close();
		System.out.println("\nDone!\n");
	}

	public void makeChrT(){
		try{
			//for each gene model
			for (int i=0; i< models.length; i++){

				//fetch exons
				ExonIntron[] exons = models[i].getExons();
				//make name and strand
				String strand = models[i].getStrand();
				String name = models[i].getDisplayName();
				if (name== null) name = models[i].getName();
				StringBuilder seqSB = new StringBuilder();

				//for each exon
				for (int j=0; j< exons.length; j++){
					//add exon sequence
					String exonSeq = new String (seq.substring(exons[j].getStart(), exons[j].getEnd()));
					seqSB.append(exonSeq);
				}

				//define start stop
				int start = base;
				int stop = start+ seqSB.length();

				//flip seq?
				String transcript;
				if (strand.equals("-")) transcript = Seq.reverseComplementDNA(seqSB.toString());
				else transcript = seqSB.toString();

				//add
				outSeq.print(transcript);
				outSeq.println("NNNNNNNNNNNNNNNNNNNN");
				base += (transcript.length()+ 20);

				//print bed
				outBed.println("chrT\t"+start+"\t"+stop+"\t"+name+"\t"+(stop-start)+"\t"+strand);
			}

		} catch (Exception e){
			e.printStackTrace();
		}
	}

	public void makeExons(){
		try{
			for (int i=0; i< models.length; i++){
				//make StringBuilder to hold sequence
				StringBuilder compositeSeq = new StringBuilder();
				//fetch exons
				ExonIntron[] exons = models[i].getExons();
				//make name and strand
				String strand = models[i].getStrand();
				String name = models[i].getDisplayName();
				if (name== null) name = models[i].getName();
				if (appendCoordinates) name = name+"_"+chr+":"+models[i].getTxStart()+"-"+models[i].getTxEnd();

				//for each exon
				for (int j=0; j< exons.length; j++){
					//define the end
					int end = exons[j].getEnd();

					//add exon sequence
					int exSubStart = exons[j].getStart() - sequenceLengthRadius;
					if (exSubStart < 0) exSubStart = 0;
					int exAddEnd = exons[j].getEnd() + sequenceLengthRadius;
					if (exAddEnd > seq.length()) exAddEnd = seq.length();
					String exonSeq = new String (seq.substring(exSubStart, exAddEnd));
					if (strand.equals("-")) exonSeq = Seq.reverseComplementDNA(exonSeq);
					compositeSeq.append(exonSeq);
					if ((j+1) != exons.length) compositeSeq.append("NNNNNNNNNN");

					//add junctions
					for (int k=j+1; k< exons.length; k++){
						int start = exons[k].getStart();

						//make seqFusion
						String seqFus = seq.substring(end-sequenceLengthRadius, end) +seq.substring(start, start+sequenceLengthRadius);
						if (strand.equals("-")) seqFus = Seq.reverseComplementDNA(seqFus);

						//print
						compositeSeq.append(seqFus);
						compositeSeq.append("NNNNNNNNNN");
					}
				}
				//make name
				outSeq.println(">"+name+"\n"+compositeSeq);
			}
		} catch (Exception e){
			e.printStackTrace();
		}
	}


	public static void main(String[] args) {
		if (args.length ==0){
			printDocs();
			System.exit(0);
		}
		new MakeTranscriptome(args);
	}		

	/**This method will process each argument and assign new varibles*/
	public void processArgs(String[] args){
		File[] fastas = null;
		File ucscTableFile = null;
		Pattern pat = Pattern.compile("-[a-z]");
		System.out.println("\nArguments: "+Misc.stringArrayToString(args, " ")+"\n");
		for (int i = 0; i<args.length; i++){
			String lcArg = args[i].toLowerCase();
			Matcher mat = pat.matcher(lcArg);
			if (mat.matches()){
				char test = args[i].charAt(1);
				try{
					switch (test){
					case 'f': fastas = IO.extractFiles(new File (args[++i]), ".fasta"); break;
					case 's': sequenceLengthRadius = Integer.parseInt(args[++i]); break;
					case 'u': ucscTableFile = new File (args[++i]); break;
					case 'r': resultsFastaFile = new File (args[++i]); break;
					case 'a': appendCoordinates = true; break;
					case 't': makeTranscriptChromosome = true; break;
					case 'h': printDocs(); System.exit(0);
					default: Misc.printExit("\nProblem, unknown option! " + mat.group());
					}
				}
				catch (Exception e){
					Misc.printExit("\nSorry, something doesn't look right with this parameter: -"+test+"\n");
				}
			}
		}
		//check params
		if (makeTranscriptChromosome == false && sequenceLengthRadius == 0) Misc.printExit("\nPlease enter a sequence length radius.\n");
		if (ucscTableFile == null || ucscTableFile.canRead() == false) Misc.printExit("\nCannot find or read your UCSC gene table?!\n");
		if (fastas == null || fastas.length ==0) Misc.printExit("\nCannot find your chromosome specific xxx.fasta files?\n");
		if (resultsFastaFile == null) Misc.printExit("\nPlease provide a full path file name for saving the spliced fasta results file.\n");
		//load UCSC table
		UCSCGeneModelTableReader reader = new UCSCGeneModelTableReader(ucscTableFile,0);
		chromSpecificGeneLines = reader.getChromSpecificGeneLines();
		//load fastaFiles
		fastaFiles = new HashMap<String,File>();
		for (int i=0; i< fastas.length; i++){
			String chr = fastas[i].getName().replaceFirst(".fasta", "");
			fastaFiles.put(chr, fastas[i]);
		}
		//check that all chroms are there in the fasta directory
		Iterator<String> it = chromSpecificGeneLines.keySet().iterator();
		while (it.hasNext()){
			String chr = it.next();
			if (fastaFiles.containsKey(chr) == false) Misc.printExit("\nCould not find fasta sequence for "+chr+"! Either delete the offending lines from " +
			"the UCSC gene table or put an appropriate chromosme in the fasta directory.\n");
		}
	}	

	public static void printDocs(){
		System.out.println("\n" +
				"**************************************************************************************\n" +
				"**                          Make Transcriptome: March 2009                          **\n" +
				"**************************************************************************************\n" +
				"MT creates a multi fasta file containing sequences representing all possible linear\n" +
				"splice junctions and extended exons for each gene in the UCSC gene table. The size of\n" +
				"each junction is 2x the sequence length radius, the exon extension in 1x on each end.\n" +
				"Be very careful about the coordinate system used in the gene table to define the\n" +
				"start and stop of exons.  UCSC uses interbase and this is assumed in this app.\n" +
				"Sequences returned are the sense + strand. Note, ELAND does not like '.' in the fasta\n" +
				"headers so be sure these are not part of your gene names.\n\n" +

				"Options:\n"+
				"-f Fasta file directory, should contain chromosome specific xxx.fasta files.\n" +
				"-u UCSC gene table, full path. See, http://genome.ucsc.edu/cgi-bin/hgTables\n"+
				"-s Sequence length radius.\n" +
				"-r Results fasta file, full path.\n"+
				"-a Append gene coordinates onto name.\n"+
				"-t Make transcript chromosome 'chrT'.\n"+

				"\n"+

				"Example: java -Xmx1500M -jar pathTo/USeq/Apps/MakeTranscriptome -s 36 \n" +
				"      -f /Genomes/Hg18/Fastas/ -u /Anno/Hg18/ucscKnownGenes.txt -r\n" +
				"      /Genomes/Hg18/Fastas/hg18_32_splices.fasta \n\n" +

		"************************************************************************************\n");

	}
}
